package org.anno.chunkdb.backend;

import org.anno.chunkdb.Chunkdb;
import org.anno.chunkdb.backend.query.StringResultColumn;

import java.io.IOException;
import java.io.OutputStream;
import java.util.Comparator;
import java.util.Iterator;
import java.util.function.Consumer;
import java.util.function.ToIntFunction;
import java.util.stream.IntStream;
import java.util.stream.Stream;

/**
 * Stores several strings as concatenated UTF-8 bytes with split points.
 */
public class StringColumn extends IndexedColumn implements Table.Column<FlyRow>,
                                                           StringResultColumn<FlyRow> {

  private final byte[] data;
  private final int[] split;

  private StringColumn(byte[] data, int[] split, int[] index) {
    super(index);
    this.data = data;
    this.split = split;
  }

  public Comparator<FlyRow> getComparator() {
    return Comparator.comparing(this::getString);
  }

  public ToIntFunction<FlyRow> getDiffFunction(Chunkdb.Filter filter) {
    byte[] data = filter.getStringRhsBytes().toByteArray();
    ByteSlice rhs = new ByteSlice(data, 0, data.length);
    return r->getString(r).compareTo(rhs);
  }

  public ByteSlice getString(FlyRow row) {
    return new ByteSlice(data, split[row.i], split[row.i + 1]);
  }

  /**
   * Returns smallest i s.t. v[index[i]] >= min
   */
  public int getGERank(byte[] min) {
    ByteSlice minSlice = new ByteSlice(min, 0, min.length);
    int start = 0;
    int end = index.length;
    while (start < end) {
      int mid = (start + end) / 2;
      int i = index[mid];
      if (new ByteSlice(data, split[i], split[i + 1]).compareTo(minSlice) < 0) {
        start = mid + 1;
      } else {
        end = mid;
      }
    }
    return start;
  }

  /**
   * Returns largest i s.t. i == -1 || v[index[i]] < v. Example: rank("x",{"x","x"}) == -1,
   * rank("xa", {"x","x"}) == 1.
   */
  public int getLTRank(byte[] max) {
    ByteSlice maxSlice = new ByteSlice(max, 0, max.length);
    int start = -1;
    int end = index.length - 1;
    while (start < end) {
      int mid = (start + end + 1) / 2;
      int i = index[mid];
      if (new ByteSlice(data, split[i], split[i + 1]).compareTo(maxSlice) < 0) {
        start = mid;
      } else {
        end = mid - 1;
      }
    }
    return start;
  }

  public void writeValue(FlyRow flyRow, OutputStream out) throws IOException {
    int i = flyRow.i;
    int start = split[i];
    out.write(data, start, split[i + 1] - start);
  }

  public static Builder builder() {
    return new Builder();
  }

  public static class Builder implements Consumer<byte[]> {

    private int size;
    private int lengthSum;
    private final IntStream.Builder splits = IntStream.builder();
    private final Stream.Builder<byte[]> strings = Stream.builder();

    private Builder() {
      splits.accept(0);
    }

    public void accept(byte[] bytes) {
      size += 1;
      lengthSum += bytes.length;
      splits.accept(lengthSum);
      strings.accept(bytes);
    }

    public StringColumn build() {
      byte[] data = new byte[lengthSum];
      int i = 0;
      for (Iterator<byte[]> it = strings.build().iterator(); it.hasNext();) {
        byte[] bytes = it.next();
        System.arraycopy(bytes, 0, data, i, bytes.length);
        i += bytes.length;
      }

      int[] split = splits.build().toArray();
      int[] index = IntStream.range(0, size)
          .mapToObj(FlyRow::new)
          .sorted(Comparator.comparing(r -> new ByteSlice(data, split[r.i], split[r.i + 1])))
          .mapToInt(r -> r.i).toArray();
      return new StringColumn(data, split, index);
    }
  }
}
